# Copyright 2025 Norris
# MIT-0
#
'''
Loads a data set into OpenSearch Service to serve as a knowledge base for
RAG retrieval. This code uses the OpenSearch Python client, provided by  
opensearch-py, to call OpenSearch's _bulk API with the data. 

It creates embeddings for the documents to support semantic search for
the retrieval via an ingest pipeline. 
'''

from opensearchpy import OpenSearch
import os


opensearch_port = 443
opensearch_service_api_endpoint = os.environ['OPENSEARCH_SERVICE_DOMAIN_ENDPOINT']
opensearch_user_name = os.environ['OPENSEARCH_SERVICE_ADMIN_USER']
opensearch_user_password = os.environ['OPENSEARCH_SERVICE_ADMIN_PASSWORD']
embedding_model_id = os.environ['EMBEDDING_MODEL_ID']
index_name = "opensearch_kl_index"


# Ensure the endpoint matches the contract for the opensearch-py client. Endpoints
# are specified without the leading URL scheme or trailing slashes.
if opensearch_service_api_endpoint.startswith('https://'):
  opensearch_service_api_endpoint = opensearch_service_api_endpoint[len('https://'):]
if opensearch_service_api_endpoint.endswith('/'):
  opensearch_service_api_endpoint = opensearch_service_api_endpoint[:-1]


# The mapping sets kNN to true to enable vector search for the index. It defines
# the text field as type text, and a text_embedding field that uses the FAISS engine
# for storage and retrieval, using the HNSW algorithm.
mapping = {
    "settings": {
        "index": {
            "knn": True,
            "number_of_shards": 1,
            "number_of_replicas": 2
        }
    },
    "mappings": {
        "properties": {
            "text": {"type": "text"},
            "text_embedding": {
              "type": "knn_vector",
              "dimension": 1024,
              "method": {
                  "name": "hnsw",
                  "space_type": "l2",
                  "engine": "faiss",
                  "parameters": {"ef_construction": 128, "m": 24},
              }
            }
        }
    }
}


# The data for the knowledge base.
population_data = [
  {"index": {"_index": index_name, "_id": "1"}},
  {"text": "Amazon OpenSearch Service 是一项托管服务，可以轻松地在 AWS 云中部署、操作和扩展 OpenSearch 集群。 OpenSearch 服务域是 OpenSearch 集群的同义词。域是包含您指定的设置、实例类型、实例计数和存储资源的集群。亚马逊 OpenSearch 服务支持 OpenSearch 传统的 Elasticsearch OSS（最高 7.10，即该软件的最终开源版本）。创建域时，您可以选择使用哪种搜索引擎。"},
  {"index": {"_index": index_name, "_id": "2"}},
  {"text": "OpenSearch是一个完全开源的搜索和分析引擎，用于日志分析、实时应用程序监控和点击流分析等用例。"},
  {"index": {"_index": index_name, "_id": "3"}},
  {"text": "Amazon OpenSearch 服务的特点\\n大量 CPU、内存和存储容量配置，也称为实例类型，包括具有成本效益的 Graviton 实例。\\n支持多达 1002 个数据节点\\n高达 25 PB 的连接存储\\n为只读数据提供经济实惠UltraWarm的冷存储"},
  {"index": {"_index": index_name, "_id": "4"}},
  {"text": """OpenSearch 与亚马逊 OpenSearch 服务相比，何时使用该服务\\n对于开源OpenSearch 您的组织愿意手动监控和维护自行预置的集群，并且拥有具备相应技能的人员。
您需要对代码拥有完全的编译级别控制。
您的组织希望或非常独特地使用开源软件。
您执行多云战略，需要不特定于供应商的技术。
您的团队有能力解决任何关键的生产问题。
您希望能够灵活地以任何需要的方式使用、修改和扩展产品。
您希望在新功能发布后立即使用这些功能。"""},
  {"index": {"_index": index_name, "_id": "5"}},
  {"text": """OpenSearch 与亚马逊 OpenSearch 服务相比，何时使用该服务，对于Amazon OpenSearch托管服务，您不想手动管理、监控和维护基础设施。
您想利用 Amazon S3 的持久性和低成本优势，通过跨存储层进行数据分层，从而管理不断增加的分析成本。
你想利用与其他数据库的集成， AWS 服务 例如 DynamoDB、Amazon DocumentDB（兼容 MongoDB）、IAM 和。 CloudWatch CloudFormation
在预防性维护和生产期间出现问题时，您需要轻松获得帮助。 支持
您想利用自我修复、主动维护、韧性和备份等功能。"""},
  {"index": {"_index": index_name, "_id": "6"}},
  {"text": "Amazon OpenSearch Serverless 是 Amazon OpenSearch 服务的按需无服务器选项，它消除了配置、配置和调整 OpenSearch 集群的操作复杂性。它非常适合那些不愿自行管理集群或缺乏用于大规模部署的专用资源和专业知识的组织。借助 OpenSearch Serverless，您可以搜索和分析大量数据，而无需管理底层基础架构。"},
    {"index": {"_index": index_name, "_id": "7"}},
  {"text": "OpenSearch Serverless 集合是一组 OpenSearch 索引，它们协同工作以支持特定的工作负载或用例。与需要手动配置的自我管理 OpenSearch 集群相比，集合简化了操作。"},
    {"index": {"_index": index_name, "_id": "8"}},
  {"text": "OpenSearch Serverless 支持与 OpenSearch 开源套件相同的采集和查询 API 操作，因此您可以继续使用现有的客户端和应用程序。您的客户端必须与 OpenSearch 2.x 兼容，才能使用 OpenSearch Serverless。"},
    {"index": {"_index": index_name, "_id": "9"}},
  {"text": "Amazon OpenSearch Ingestion 是一个完全托管的无服务器数据收集器，可将实时日志、指标和跟踪数据流式传输到亚马逊 OpenSearch 服务域和 OpenSearch 无服务器集合。"},
    {"index": {"_index": index_name, "_id": "10"}},
    {"text": "借助 OpenSearch Ingestion，您不再需要像 Logstash 或 Jaeger 这样的第三方工具来摄取数据。您可以将数据生成器配置为将数据发送到 OpenSearch Ingestion，然后它会自动将其传输到您的指定域或集合。您也可以在交付数据之前转换数据。"}
    {"index": {"_index": index_name, "_id": "11"}},
    {"text": "UltraWarm 为在 Amazon OpenSearch 服务上存储大量只读数据提供了一种经济实惠的方式。标准数据节点使用“热”存储，其形式是连接到每个节点的实例存储或 Amazon EBS 卷。热存储为编制索引和搜索新数据提供尽可能快的性能。"}
    {"index": {"_index": index_name, "_id": "11"}},
    {"text": "UltraWarm 节点不使用附加存储，而是使用 Amazon S3 和复杂的缓存解决方案来提高性能。对于不主动写入、查询频率较低且不需要相同性能的索引， UltraWarm 可以显著降低每 GiB 数据的成本。因为除非将温索引返回到热存储，否则它们 UltraWarm 是只读的，因此最适合存储不可变的数据，例如日志。"}
    {"index": {"_index": index_name, "_id": "11"}},
    {"text": "UltraWarm 节点不使用附加存储，而是使用 Amazon S3 和复杂的缓存解决方案来提高性能。对于不主动写入、查询频率较低且不需要相同性能的索引， UltraWarm 可以显著降低每 GiB 数据的成本。因为除非将温索引返回到热存储，否则它们 UltraWarm 是只读的，因此最适合存储不可变的数据，例如日志。"}
]


# OpenSearch ingest pipelines let you define processors to apply to your documents at ingest
# time. The text_embedding processor lets you select a source field and a destination field
# for the embedding. At ingest, OpenSearch will call the model, via its model id, to 
# generate the embedding.
ingest_pipeline_definition = {
  "processors": [
    {
      "text_embedding": {
        "model_id": embedding_model_id,
        "field_map": {
            "text": "text_embedding"
        }
      }
    }
  ]
}


# Set up for the client to call OpenSearch Service
hosts = [{"host": opensearch_service_api_endpoint, "port": opensearch_port}]
client = OpenSearch(
    hosts=hosts,
    http_auth=(opensearch_user_name, opensearch_user_password),
    use_ssl=True,
    verify_certs=False,
    ssl_assert_hostname=False,
    ssl_show_warn=False,
)


# Check whether an index already exists with the chosen name. If you receive an 
# exception, change the index name in the global variable above, and be sure 
# also to change the index_name global variable in run_rag.py
if client.indices.exists(index=index_name) == 200:
  raise Exception(f'Index {index_name} already exists. Please choose a different name in load_data.py. Be sure to change the index_name in run_rag.py as well.')

# This code does not validate the response. In actual use, you should wrap this
# block in try/except and validate the response. 
client.indices.create(index=index_name, body=mapping)
r = client.ingest.put_pipeline(id="embedding_pipeline", 
                               body=ingest_pipeline_definition)
client.bulk(index=index_name, 
            body=population_data,
            pipeline="embedding_pipeline", 
            refresh=True)

print(f'Loaded data into {index_name}')